/****************************************************************************
 *
 * Copyright (C) 2003-2010 Los Alamos National Security, LLC
 *                         Packet Analytics Corporation
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License Version 2 as
 * published by the Free Software Foundation.  You may not use, modify or
 * distribute this program under any other version of the GNU General
 * Public License.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 ****************************************************************************/
package nfse.mongo;

import com.mongodb.Mongo;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.BasicDBObject;
import com.mongodb.DBObject;
import com.mongodb.DBCursor;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.net.URL;
import java.net.URLConnection;
import java.util.ArrayList;
import java.util.Random;

import org.bson.BasicBSONObject;

import nfse.DataFileProcessor;
import nfse.NetFSE;
import nfse.NetFSERecord;
import nfse.NetFSEUtil;
import nfse.SensorInfo;
import nfse.ToolInfo;
import nfse.gui.record.FlexRecord;
import nfse.gui.record.FlexRecordField;
import nfse.gui.server.GUIRecordFactory;

public class MongoStore extends DataFileProcessor {

    @Override
    public void process() throws Exception {
        /**
         * Behind the scenes, the system tracks how much of a file has been
         * processed if any. Call getOffset() to skip to the next record to be
         * processed. If it is non-zero, moveToOffset() will move to the correct
         * location in the file.
         */
        int offset = getOffset();
        if (offset > 0) {
            moveToOffset();
        }

        /**
         * Need to know what kind of data we are dealing with. The "Tool" is a
         * legacy naming convention that should really be "Type". The Sensor is
         * an instance of the Tool/Type. For example you would likely have
         * multiple Snort sensors. They share a common Tool but are different
         * Sensors.
         */
        SensorInfo sensorInfo = NetFSE.getSensorInfo(this.getSensorID());
        ToolInfo toolInfo = NetFSE.getToolInfo(sensorInfo.getToolID());

        /**
         * This is the class needed to parse records. Determined by the config
         * file and system metadata.
         */
        GUIRecordFactory rec = (GUIRecordFactory) NetFSE.getRecordFactory(toolInfo.getId());
        FlexRecordField[] columns = rec.getFlexRecordFields();
        Mongo m = new Mongo("localhost");

        DB db = m.getDB("test");
        DBCollection coll = db.getCollection("foo");

        ArrayList<DBObject> docs = new ArrayList<DBObject>();

        long startTS = System.currentTimeMillis();
        long count = 0;
        byte[] bytes = null;
        while ((bytes = getNextRecord()) != null) {
            NetFSERecord temp = null;
            try {
                temp = rec.generate(new String(bytes), this.getSensorID());
            } catch (Exception e) {
                e.printStackTrace();
                rec = null;
            }

            if (temp != null) {
                FlexRecord fr = rec.getFlexRecord(temp);
                BasicDBObject doc = new BasicDBObject();
                for (int i = 0; i < fr.values.length; i++) {
                    if (columns[i].type == FlexRecordField.TYPE_STRING
                            && columns[i].subType == FlexRecordField.SUB_TYPE_IP_ADDR) {
                        doc.put(columns[i].columnName, NetFSEUtil.convertIPtoLong(fr.values[i]));
                    } else if (columns[i].type == FlexRecordField.TYPE_INT) {
                        doc.put(columns[i].columnName, Integer.parseInt(fr.values[i]));
                    } else if (columns[i].type == FlexRecordField.TYPE_LONG) {
                        doc.put(columns[i].columnName, Long.parseLong(fr.values[i]));
                    } else if (columns[i].type == FlexRecordField.TYPE_FLOAT) {
                        doc.put(columns[i].columnName, Double.parseDouble(fr.values[i]));
                    } else {
                        doc.put(columns[i].columnName, fr.values[i]);
                    }
                }
                doc.put("netfse_record", temp.getInitialString());
                docs.add(doc);
                count++;
                if (count % 1000 == 0) {
                    coll.insert(docs);
                    docs.clear();
                }
            }
        }

        if (docs.size() > 0) {
            coll.insert(docs);
        }

        m.close();

        double sec = (System.currentTimeMillis() - startTS) / 1000.0;
        if (count > 0)
            System.out.println("Count=" + count + ", RPS: " + count / sec);
    }

}
